#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def Gen_BAT_File(id):
    BAT_TEMPLATE = """

@echo off
set "msg=download failed, retry"

rem running aspera downloader
echo start downloading
".\Aspera\Aspera Connect\\bin\\ascp" -i ".\Aspera\Aspera Connect\etc\\asperaweb_id_dsa.openssh" -k1 -QTr -l180m --user=anonftp --host=ftp.ncbi.nlm.nih.gov --mode=recv --file-pair-list={file_pair_list}.lst --file-manifest=text --file-manifest-path=log_mani -L log .

if "%errorlevel%"=="0" ( set "msg=download finished")

echo %msg%

mshta vbscript:msgbox(%0"%msg%",64,%0"download")(window.close)

"""

    bat_file = open("D:/public_data_download/" + id + '.bat', 'w', encoding='GBK')
    bat_file.write(BAT_TEMPLATE.format(file_pair_list=id))
    bat_file.close()


def Gen_List_File(id):
    client_datalist, collection_datalist = InitMongoDB("10.10.116.28", 27017, "ant_nest", "datalist")

    path_file = open("D:/public_data_download/" + id + ".lst", 'w', encoding='utf-8')

    for result in collection_datalist.find({"JobID": id, "data_store_path": {"$exists": 0}}, {"_id": 0, "RunID": 1}):
        if "RunID" in result:
            run_id = result["RunID"]
            path_file.write('/sra/sra-instant/reads/ByRun/sra/{}/{}/{}/{}.sra\n'.format(
                run_id[:3], run_id[:6], run_id, run_id))
            path_file.write('{}/{}.sra\n'.format(id, run_id))
            path_file.flush()
    path_file.close()

    CloseMongoDB(client_datalist)


if __name__ == '__main__':
    for i in range(16, 31):
        id = "dlst" + str(i).zfill(6)
        Gen_BAT_File(id)
        Gen_List_File(id)
